import os
import logging
from typing import Optional
from contextlib import contextmanager
from chromadb import PersistentClient
from pydantic_settings import BaseSettings
from pydantic import Field
from dotenv import load_dotenv

# 初始化日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class Config(BaseSettings):
    """应用配置类"""
    deepseek_api_key: str = os.getenv("DEEPSEEK_API_KEY")
    db_host: str = os.getenv("DB_HOST")
    db_port: int = os.getenv("DB_PORT")
    db_user: str = os.getenv("DB_USER")
    db_password: str = os.getenv("DB_PASSWORD")
    db_name: str = os.getenv("DB_PASSWORD")
    
    class Config:
        env_file = ".env"
        env_file_encoding = "utf-8"

config = Config()

class DatabaseManager:
    """数据库连接管理器"""
    def __init__(self):
        self._conn_pool = []
    
    @contextmanager
    def get_connection(self):
        """获取数据库连接的上下文管理器"""
        conn = self._create_connection()
        try:
            yield conn
        finally:
            self._release_connection(conn)
    
    def _create_connection(self):
        """创建新数据库连接"""
        # 实际实现需要根据数据库类型调整
        import mysql.connector
        return mysql.connector.connect(
            host=config.db_host,
            port=config.db_port,
            user=config.db_user,
            password=config.db_password,
            database=config.db_name
        )
    
    def _release_connection(self, conn):
        """释放连接到连接池"""
        if conn.is_connected():
            self._conn_pool.append(conn)

class DeepSeekVanna:
    def __init__(self, db_manager: DatabaseManager):
        self.db_manager = db_manager
        self.api_key = config.deepseek_api_key
        self.chroma_client = PersistentClient(path="./chroma_data")
        
    def generate_sql(self, question: str, chat_history: list = None, **kwargs) -> str:
        """生成SQL查询语句"""
        import requests
        from mysql.connector import Error as DBError
        
        # 构建包含历史记录的上下文
        messages = [{"role": "system", "content": "你是一个专业的SQL工程师，根据以下数据库结构和历史示例生成SQL"}]
        
        if chat_history:
            messages.extend(chat_history[-3:])  # 只保留最近3条历史
            
        with self.db_manager.get_connection() as conn:
            schema = self._get_schema(conn)
            messages.append({"role": "user", "content": f"问题：{question}\n表结构：{schema}"})
        
        try:
            # 初始化chromadb集合
            collection = self.chroma_client.get_or_create_collection(
                name="sql_queries",
                metadata={"hnsw:space": "cosine"}
            )
            
            # 搜索相似的历史查询
            results = collection.query(
                query_texts=[question],
                n_results=3
            )
            
            # 构建历史查询上下文
            history_context = ""
            if results['documents']:
                history_context = "\n历史相似查询示例：\n" + "\n".join(
                    f"- {query}" for query in results['documents'][0]
                )
            
            with self.db_manager.get_connection() as conn:
                schema = self._get_schema(conn)
                
                # 获取相关文档上下文
                doc_collection = self.chroma_client.get_collection("documentation")
                doc_results = doc_collection.query(
                    query_texts=[question],
                    n_results=2
                )
                doc_context = ""
                if doc_results['documents']:
                    doc_context = "\n相关文档：\n" + "\n".join(
                        f"- {doc}" for doc in doc_results['documents'][0]
                    )

                headers = {
                    "Authorization": f"Bearer {self.api_key}",
                    "Content-Type": "application/json"
                }
                payload = {
                    "messages": [
                        {"role": "system", "content": "你是一个专业的SQL工程师，根据以下数据库结构、历史示例和文本规则生成SQL，避免数据库结构以外的表结构、表名出现"},
                        {"role": "user", "content": f"问题：{question}\n表结构：{schema}\n历史示例:{history_context}\n数据规则：{doc_context}"}
                    ],
                    "model": "deepseek-chat",
                    "temperature": 1.0
                }
                
                logger.info(f"Sending request to DeepSeek API: {payload}")
                response = requests.post(
                    os.getenv("DEEPSEEK_URL"),
                    headers=headers,
                    json=payload,
                    timeout=100
                )
                # 检查响应状态
                if response.status_code != 200:
                    logger.error(f"API响应异常: {response.status_code} - {response.text}")
                    raise Exception(f"API返回异常状态码: {response.status_code}")

                # 解析JSON前检查内容
                try:
                    response_data = response.json()
                except Exception as e:
                    logger.error(f"JSON解析失败: {str(e)} 原始响应: {response.text}")
                    raise Exception("API返回了无效的JSON数据")

                # 检查关键字段是否存在
                if 'choices' not in response_data or len(response_data['choices']) == 0:
                    logger.error(f"API响应缺少必要字段: {response_data}")
                    raise Exception("API响应结构异常")
                
                result = response_data['choices'][0]['message']['content']
                logger.debug(f"Received SQL response: {result}")
                
                # 存储新的查询到向量数据库
                collection.add(
                    documents=[question],
                    metadatas=[{"sql": result}],
                    ids=[str(hash(question))]
                )
                logger.info(f"返回的结果为：{result}")
                return result
                
        except requests.exceptions.RequestException as e:
            logger.error(f"API请求失败: {str(e)}")
            raise Exception("API服务暂时不可用") from e
        except DBError as e:
            logger.error(f"数据库操作失败: {str(e)}")
            raise Exception("数据库连接异常") from e
        except KeyError as e:
            logger.error(f"API响应格式异常: {str(e)}")
            raise Exception("API响应解析失败") from e
    
    def train(self, ddl: str = None, sql: str = None, question: str = None, documentation: str = None):
        """存储训练数据到向量数据库"""
        # 处理文档类型训练数据
        if documentation:
            collection = self.chroma_client.get_or_create_collection(
                name="documentation",
                metadata={"hnsw:space": "cosine"}
            )
            collection.add(
                documents=[documentation],
                ids=[str(hash(documentation))]
            )
            return
        if ddl:
            collection = self.chroma_client.get_or_create_collection(
                name="ddl_schema",
                metadata={"hnsw:space": "cosine"}
            )
            collection.add(
                documents=[ddl],
                ids=[str(hash(ddl))]
            )
            
        if sql and question:
            collection = self.chroma_client.get_or_create_collection(
                name="sql_examples",
                metadata={"hnsw:space": "cosine"}
            )
            collection.add(
                documents=[question],
                metadatas=[{"sql": sql}],
                ids=[str(hash(question))]
            )

    def _get_schema(self, conn) -> str:
        """获取数据库表结构"""
        from mysql.connector import Error as DBError
        
        try:
            with conn.cursor() as cursor:
                cursor.execute("SHOW TABLES")
                tables = [row[0] for row in cursor.fetchall()]
                schema = []
                
                for table in tables:
                    try:
                        cursor.execute(f"SHOW CREATE TABLE {table}")
                        schema.append(cursor.fetchone()[1])
                    except DBError as e:
                        logger.warning(f"获取表 {table} 结构失败: {str(e)}")
                        
                return "\n".join(schema)
        except DBError as e:
            logger.error(f"获取数据库元数据失败: {str(e)}")
            raise
